///////////////////////////////////////////
// Authors : Mathieu Robitaille - 6396682
//           Sacha Sukhdeo - 6054412
// Class : COMP445
//         Assignment 2
///////////////////////////////////////////

#include "TcpBase.h"

#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <vector>

TIMEVAL TcpBase::TIMEOUT = {1l, 300000l};
const int TcpBase::WINDOW_SIZE = 2;
const int TcpBase::SEQ_NB_RANGE = 2*TcpBase::WINDOW_SIZE + 1;
const int TcpBase::LAST_PACKET_MAX_RESEND = 5;


unsigned long TcpBase::ResolveName(char name[]) {
	OutputDebugStringW(L"\n\nTcpBase::ResolveName\n\n");
    struct hostent *host;            /* Structure containing host information */

    if ((host = gethostbyname(name)) == NULL)
        err_sys("gethostbyname() failed");

    /* Return the binary, network byte ordered address */
    return *((unsigned long *) host->h_addr_list[0]);
}

/*
msg_recv returns the length of bytes in the msg_ptr->buffer,which have been recevied successfully.
*/
int TcpBase::msg_recv(SOCKET s, Msg * p_msg, SOCKADDR* p_from) {
    int sizeFromAddr = sizeof(SOCKADDR_IN);
    int n;
    if ((n = recvfrom(s, (char*)p_msg, sizeof(Msg), 0, (SOCKADDR*)p_from, p_from != NULL ? &sizeFromAddr : NULL)) <= 0)
        err_sys("Receive message error. Error : %d", WSAGetLastError());

    return n;
}

 int TcpBase::msg_send(SOCKET s, Msg * p_msg, SOCKADDR* p_addr, int addrSize)
{
    unsigned int n;
	if(p_msg->length < 0 || p_msg->length > BUFFER_LENGTH)
		std::cerr << "Invalid buffer length" << std::endl;
	p_msg->length = min(max(p_msg->length, 0), BUFFER_LENGTH);
    for (n = 0; n < MSGHDRSIZE+p_msg->length;) {
        n += sendto(s, (char*) p_msg, MSGHDRSIZE+p_msg->length, 0, p_addr, addrSize);
        if (n == SOCKET_ERROR)
            err_sys("Error while sending the message. Error : %d", WSAGetLastError());
    }
    return n;
}

int TcpBase::WaitForSocket(FD_SET fds, LPTIMEVAL timeout) {
    int n = select(0, &fds, NULL, NULL, timeout);
    if (n == SOCKET_ERROR)
        err_sys("Error while waiting for a socket to be ready. Error : %d", WSAGetLastError());
    return n;
}

void TcpBase::err_sys(char * fmt,...)
{     
    perror(NULL);
    va_list args;
    va_start(args,fmt);
    fprintf(stderr,"error: ");
    vfprintf(stderr,fmt,args);
    fprintf(stderr,"\n");
    va_end(args);
    exit(1);
}

void TcpBase::SendFile(const std::string& filename, SOCKET sTo, SOCKET sFrom, const std::string& hostname, SOCKADDR* p_addr, int addrSize) {
    std::ifstream ifs(filename, std::ifstream::binary);
    if (!ifs) {
        std::cerr << "File '" << filename << "' does not exists." << std::endl;
		Msg msg;
		msg.type = MSG_ERR;
		char m[] = "file not found";
		memcpy(msg.buffer, m,sizeof(m));
		memset(msg.buffer+sizeof(m), 0,1);
		msg.length = sizeof(m);
		//msg_send(s, &msg, p_addr, addrSize);
        LogInfo li;
        memset(&li, 0, sizeof(LogInfo));
        SendPacket(sTo, sFrom, msg, p_addr, addrSize, li, true);
        return;
	}

    Log("Sender : Starting on host [%s]", m_hostname);
    Log("Sender : Sending file '%s'", filename.c_str());

	std::cout << "sending file to " << hostname << " waiting ..." << std::endl; 
    ifs.seekg(0, ifs.end);
    int length = (int)ifs.tellg() + 1; // Make sure to read the EOF otherwise we have an infinite loop
    ifs.seekg(0, ifs.beg);
    int i = 1;
    int count = (int)std::ceilf((float)length/FILEDATA_BUFFER_LENGTH);

    LogInfo li;
    memset(&li, 0, sizeof(LogInfo));
    bool showFirstPacket = true;

    while (ifs) {
        Msg msg;
        int readLength = min(length, (int)FILEDATA_BUFFER_LENGTH);
        FileData fd;
        //fd.idx = i++;
        //fd.count = count;
        ifs.read(fd.filedata, readLength);
        if (showFirstPacket) {
            showFirstPacket = false;
            if (filename.substr(filename.length()-3, 3) == "txt") {
                std::cout << "File sample data: " << std::endl;
                std::cout.write(fd.filedata, readLength-1);
                std::cout << std::endl;
            } else {
                std::cout << "File contains binary data, no sample shown" << std::endl;
            }
        }
        fd.length = (int)ifs.gcount();
        length -= fd.length;
        msg.type = REQ_FILE;
        msg.length = fd.length+3*sizeof(int);
        msg.seqNb = localSeqNb % SEQ_NB_RANGE;
		msg.ack =0;
        memcpy(msg.buffer, &fd, msg.length);

        SendPacket(sTo, sFrom, msg, p_addr, addrSize, li);
    }
    Msg m;
    m.type = REQ_FILE;
    m.length = 0;
    memset(m.buffer, 0, sizeof(m.buffer));
    m.ack = 0;
    m.seqNb = localSeqNb;
	SendPacket(sTo, sFrom, m, p_addr ,addrSize, li, true);

    Log("Sender : File '%s' transfer completed", filename.c_str());
    Log("Sender : Number of effective bytes sent : %i", li.effectiveBytes);
	Log("Sender : Number of effective packets sent : %i", li.effectivePackets);
    Log("Sender : Number of packet sent : %i", li.packets);
    Log("Sender : Number of bytes sent : %i", li.bytes);
	std::cout << "File '" << filename << "' sent. to " << hostname << std::endl;
}

void TcpBase::ReceiveFile(const std::string& filename, SOCKET sTo, SOCKET sFrom, SOCKADDR * addrFrom, const std::string& hostname) {
    std::ofstream ofs(filename, std::ofstream::binary);
    if (!ofs) {
        std::cerr << "File '" << filename << "' could not be open for writing." << std::endl;
        return;
    }
    Log("Receiver : Starting on host [%s]", m_hostname);
    Log("Receiver : Receiving file '%s'", filename.c_str());

    Msg msg;
    FileData *pfd = NULL;
    std::cout << "Receiving file " << filename << " from " << hostname << "..." << std::endl;

    LogInfo li;
    memset(&li, 0, sizeof(LogInfo));
    do {
		msg = ReceivePacket(sFrom, sTo, addrFrom, li);

		if(msg.type == MSG_ERR) {
			std::cerr << msg.buffer << std::endl;
            return;
		}
        if (msg.type != REQ_FILE) {
            std::cerr << "Error: Wrong message type. '" << msg.type << "'." << std::endl;
			--remoteSeqNb;
			//msg = ReceivePacket(sFrom, sTo, addrFrom, li);
            // return;
        } else {
			if (msg.length > 0) {
				pfd = (FileData*)msg.buffer;
				ofs.write(pfd->filedata, pfd->length);
			}
		}
    } while(msg.length != 0);
    Log("Receiver : Transfer of file '%s' completed", filename.c_str());
	Log("Receiver : Number of effective bytes received : %i", li.effectiveBytes);
	Log("Receiver : Number of effective packets received : %i", li.effectivePackets);
    Log("Receiver : Number of bytes received : %i", li.bytes);
    Log("Receiver : Number of packet received : %i", li.packets);
    ofs.close();
    std::cout << "File " << filename << " reception is complete." << std::endl;
}

void TcpBase::SendFileList(std::string directory, SOCKET sTo, SOCKET sFrom, SOCKADDR* p_addr, int addrSize) {
    WIN32_FIND_DATA fd = { 0 };
    directory += "\\*.*";
    HANDLE hf = FindFirstFile(directory.c_str(), &fd);
    if (hf != INVALID_HANDLE_VALUE) {
        std::string fileNames = "";
        do {
            //std::cout << fd.cFileName << std::endl;
            fileNames += fd.cFileName;
            fileNames+= "\n";
        } while (FindNextFile(hf, &fd));
        Msg msg;
        memcpy(msg.buffer,fileNames.c_str(),fileNames.length()*sizeof(char));
        msg.type=RESP;
        msg.length=fileNames.length()*sizeof(char);
        LogInfo li;
        memset(&li, 0, sizeof(li));
        SendPacket(sTo, sFrom, msg, p_addr, addrSize, li, true);
    }
}

void TcpBase::ReceiveFileList(SOCKET sTo, SOCKET sFrom, SOCKADDR* addrFrom)
{
	Msg msg;
    LogInfo li;
    memset(&li, 0, sizeof(li));
	msg = ReceivePacket(sFrom, sTo, addrFrom, li);
	char *temp = new char[msg.length+1];
	memset(temp, 0, msg.length+1);
	memcpy(temp, msg.buffer, msg.length);
	std::cout << std::endl << temp << std::endl;
	delete [] temp;
}

void TcpBase::Log(char * fmt,...) {
#define TRACE
#ifdef TRACE
    m_log.open(LogFile(), std::ios_base::app|std::ios_base::out);
    va_list args;
    va_start(args, fmt);
    char s[256];
    memset(s, 0, sizeof(s));
    vsprintf_s(s, sizeof(s), fmt, args);
    OutputDebugString(s);
    OutputDebugString("\n");
    m_log << s << std::endl;
    va_end(args);
    m_log.close();
#endif // TRACE
}

// TODO add a parameter to force the sending of pending packet, if the window is not full but no more packet to send
void TcpBase::SendPacket(SOCKET sTo, SOCKET sFrom, Msg m, SOCKADDR* p_addr, int addrSize, LogInfo& li, bool isLast) {
    FD_SET fds;
    FD_ZERO(&fds);
	FD_SET(sFrom, &fds);

	bool acked = false;
	bool ackMissed = false;
	bool sendWindow = true; // should we re-send due to a missing ACK, or we got an NACK
    int resendCount = 0;
	int ackReceivedCount = 0; // used in while loop for receiving ACKS ... once ackReceivedCount reaches WINDOW_SIZE...we stop looping,
							  // ...in the case of a received NACK ackReceivedCount represents the index value where the NACK'D Msg is stored in vector<Msg>msgBuffer.
	std::vector<Msg> tempMsgVector;
	tempMsgVector = msgBuffer;
    // You could use #include <vector> with std::vector<Msg> as a buffer
    
	while(sendWindow) // true if we are sending for the first time OR we NEED TO RE-SEND
	{
		if(ackMissed)
		{
			for(Msg msg: tempMsgVector)
			{
				int n = msg_send(sTo, &msg, p_addr, addrSize);      
				Log("Sender : Sent packet %i", li.effectivePackets);
				li.bytes += n;
				++li.packets;
				resendCount++;
			}
			ackMissed = false;
		}
		else if(!nackReceived)
		{
			int n = msg_send(sTo, &tempMsgVector[NACKMsgIndex], p_addr, addrSize);      
				Log("Sender : Sent packet %i", li.effectivePackets);
				li.bytes += n;
				++li.packets;
				nackReceived = true;
		}
		else
		{
			while (!acked && tempMsgVector.size()< WINDOW_SIZE)
			{
				m.seqNb = localSeqNb % SEQ_NB_RANGE;
				tempMsgVector.emplace_back(m);

				int n = msg_send(sTo, &m, p_addr, addrSize);      
				Log("Sender : Sent packet %i", li.effectivePackets);

				if (resendCount == 0) {
					li.effectiveBytes += n;
					++li.effectivePackets;
				}
				li.bytes += n;
				++li.packets;
			}
		}
		
		   // don't wait on a single packet, send a full window
		while(ackReceivedCount < WINDOW_SIZE)
		{
			if (WaitForSocket(fds, &TIMEOUT))
			{
				Msg mack;
				ackReceivedCount++;
				msg_recv(sFrom, &mack);
				// TODO check for msg.type ACK or NACK
				if (mack.type == ACK && mack.ack == (localSeqNb % SEQ_NB_RANGE)) 
				{
					acked = true;
					++localSeqNb;
					Log("Sender : Received ACK for packet %i", li.effectivePackets-1);
					sendWindow = false;
				}
				else if(mack.type ==NACK)
				{
					msgBuffer = tempMsgVector;
					tempMsgVector.clear();
					sendWindow = false;
					nackReceived = false;
					NACKMsgIndex = ackReceivedCount;
					break;
				}
			}
			else
			{
				ackMissed = true;
				break;
			}
			++resendCount;
			if (isLast && resendCount >= LAST_PACKET_MAX_RESEND) {
				++localSeqNb;
				Log("Sender : Never received ACK for packet %i. Since this is the last one, assume the receiver received the packet and closed down.", li.effectivePackets);
				acked = true;
			}
			sendWindow = false;

		}
	}
}

Msg TcpBase::ReceivePacket(SOCKET sFrom, SOCKET sTo, SOCKADDR * addrFrom, LogInfo& li) {
    FD_SET fds;
    FD_ZERO(&fds);
    FD_SET(sFrom, &fds);
	Msg msg;
    bool tryAgain = true;
	while(tryAgain) {

		if ( WaitForSocket(fds, NULL)== 1) {
			Log("Receiver : Received packet %i", li.effectivePackets);
			int n = msg_recv(sFrom,&msg, addrFrom);
			li.bytes += n;
			++li.packets;

			Msg m;
			memset(&m, 0, sizeof(Msg));
            if (msg.seqNb == (remoteSeqNb % SEQ_NB_RANGE)
                || msg.seqNb == ((remoteSeqNb-1) % SEQ_NB_RANGE)) // received last packet but ack was lost, this is a retransmission
                m.type = ACK;
            else
                m.type = NACK;

			m.ack = msg.seqNb;
			msg_send(sTo, &m, addrFrom, sizeof(SOCKADDR_IN));

			Log("Receiver : Sent ACK for packet %i", li.effectivePackets);

			if (msg.seqNb == (remoteSeqNb % SEQ_NB_RANGE)) {
				tryAgain = false;
				li.effectiveBytes += n;
				++li.effectivePackets;
				++remoteSeqNb;
			}
		}
	}
    return msg;
}
